"""
数据库连接模块，提供与SQL Server数据库的连接和查询功能
"""

import os
import logging
import pandas as pd
import pyodbc
from config import DB_CONFIG

# 配置日志
logging.basicConfig(
    level=logging.INFO,
    format='%(asctime)s - %(name)s - %(levelname)s - %(message)s'
)
logger = logging.getLogger('database')

def get_connection():
    """
    创建并返回数据库连接
    """
    try:
        connection_string = (
            f"DRIVER={DB_CONFIG['driver']};"
            f"SERVER={DB_CONFIG['server']};"
            f"DATABASE={DB_CONFIG['database']};"
            f"UID={DB_CONFIG['username']};"
            f"PWD={DB_CONFIG['password']};"
            f"TrustServerCertificate=yes;"
        )
        conn = pyodbc.connect(connection_string)
        logger.info("数据库连接成功")
        return conn
    except Exception as e:
        logger.error(f"数据库连接失败: {str(e)}")
        raise

def execute_query(query, params=None):
    """
    执行SQL查询并返回结果集
    
    Args:
        query (str): SQL查询语句
        params (tuple, optional): 查询参数
        
    Returns:
        pandas.DataFrame: 查询结果集
    """
    try:
        conn = get_connection()
        logger.info(f"执行SQL查询: {query}")
        
        if params:
            df = pd.read_sql(query, conn, params=params)
        else:
            df = pd.read_sql(query, conn)
            
        conn.close()
        logger.info(f"查询完成，获取 {len(df)} 条记录")
        return df
    except Exception as e:
        logger.error(f"查询执行失败: {str(e)}")
        raise

def execute_non_query(sql, params=None):
    """
    执行非查询SQL语句（如INSERT, UPDATE, DELETE等）
    
    Args:
        sql (str): SQL语句
        params (tuple, optional): SQL参数
        
    Returns:
        int: 受影响的行数
    """
    try:
        conn = get_connection()
        cursor = conn.cursor()
        logger.info(f"执行非查询SQL: {sql}")
        
        if params:
            cursor.execute(sql, params)
        else:
            cursor.execute(sql)
            
        affected_rows = cursor.rowcount
        conn.commit()
        cursor.close()
        conn.close()
        
        logger.info(f"执行完成，影响 {affected_rows} 行")
        return affected_rows
    except Exception as e:
        logger.error(f"执行失败: {str(e)}")
        raise

def test_connection():
    """
    测试数据库连接是否正常
    """
    try:
        conn = get_connection()
        cursor = conn.cursor()
        cursor.execute("SELECT @@VERSION")
        version = cursor.fetchone()[0]
        cursor.close()
        conn.close()
        logger.info(f"数据库连接测试成功, SQL Server版本: {version}")
        return True, version
    except Exception as e:
        logger.error(f"数据库连接测试失败: {str(e)}")
        return False, str(e)

if __name__ == "__main__":
    # 测试数据库连接
    success, result = test_connection()
    if success:
        print(f"数据库连接成功! SQL Server版本: {result}")
    else:
        print(f"数据库连接失败: {result}") 